import json
import logging
import os
import re

import numpy as np
from datasets import Dataset
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from openai import OpenAI
from pymilvus import Collection, connections, utility
from ragas import evaluate
from ragas.llms import LangchainLLMWrapper
from ragas.metrics import AnswerRelevancy, Faithfulness
from sentence_transformers import SentenceTransformer

# 加载环境变量
load_dotenv()

# 设置日志
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


class Text2SQLRAGASEvaluator:
    def __init__(self):
        self.setup_models()
        self.setup_milvus()
        self.load_q2sql_data()

    def setup_models(self):
        """初始化模型"""
        # 从环境变量获取API密钥
        api_key = "sk-df3336b40bf44c0195e9bcf0fb81cdc6"

        # LLM for RAGAS
        try:
            self.llm = LangchainLLMWrapper(
                ChatOpenAI(
                    model_name="qwen-plus",
                    openai_api_key=api_key,
                    openai_api_base="https://dashscope.aliyuncs.com/compatible-mode/v1",
                )
            )
            logger.info("Qwen LLM初始化成功")
        except Exception as e:
            logger.error(f"LLM初始化失败: {e}")
            raise

        # OpenAI client for text2sql
        self.openai_client = OpenAI(
            api_key=api_key,
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
        )

        # Embedding model
        bge_m3_model = (
            "D:\\RAG\\作业2\\my_work\\finance-term-standardization\\backend\\bge-m3"
        )
        self.embedding_model = SentenceTransformer(bge_m3_model)

    def setup_milvus(self):
        """设置Milvus连接"""
        try:
            connections.connect(alias="default", host="localhost", port="19530")
            logger.info(f"成功连接到Milvus服务器: {utility.list_collections()}")
        except Exception as e:
            logger.warning(f"Milvus连接失败: {e}, 将使用模拟检索")
            self.milvus_available = False
        else:
            self.milvus_available = True

    def load_q2sql_data(self):
        """加载Q2SQL数据集"""
        try:
            with open(
                "D:\\RAG\\作业3\\rag-in-action-master\\90-文档-Data\\sakila\\q2sql_pairs.json",
                "r",
                encoding="utf-8",
            ) as f:
                self.q2sql_pairs = json.load(f)
            logger.info(f"加载了 {len(self.q2sql_pairs)} 条Q2SQL数据")
        except Exception as e:
            logger.error(f"加载Q2SQL数据失败: {e}")
            self.q2sql_pairs = []

    def local_embedding_function(self, texts):
        """本地嵌入函数"""
        if isinstance(texts, str):
            texts = [texts]
        embeddings = self.embedding_model.encode(
            texts, batch_size=12, convert_to_tensor=False, normalize_embeddings=True
        )
        return embeddings.tolist()

    def extract_entity_from_string(self, hit_string):
        """从字符串中提取实体信息"""
        entity_pattern = r"entity: ({.*})"
        entity_match = re.search(entity_pattern, hit_string, re.DOTALL)

        if entity_match:
            entity_str = entity_match.group(1)
            return eval(entity_str)
        else:
            return {}

    def retrieve(
        self, collection_name: str, query_emb: list, top_k: int = 3, fields: list = None
    ):
        """从Milvus检索相关信息"""
        if not self.milvus_available:
            return []

        try:
            collection = Collection(collection_name)
            collection.load()

            search_params = {"metric_type": "COSINE", "params": {"nprobe": 10}}

            results = collection.search(
                data=[query_emb],
                anns_field="vector",
                param=search_params,
                limit=top_k,
                output_fields=fields or [],
            )

            hits = []
            for entity in results[0]:
                hit = self.extract_entity_from_string(str(entity))
                hits.append(hit)

            return hits
        except Exception as e:
            logger.warning(f"检索失败: {e}")
            return []

    def extract_sql(self, text: str) -> str:
        """提取SQL语句"""
        # 尝试匹配 SQL 代码块
        sql_blocks = re.findall(r"```sql\n(.*?)\n```", text, re.DOTALL)
        if sql_blocks:
            return sql_blocks[0].strip()

        # 如果没有找到代码块，尝试匹配 SELECT 语句
        select_match = re.search(r"SELECT.*?;", text, re.DOTALL | re.IGNORECASE)
        if select_match:
            return select_match.group(0).strip()

        # 尝试匹配其他SQL语句
        sql_match = re.search(
            r"(INSERT|UPDATE|DELETE).*?;", text, re.DOTALL | re.IGNORECASE
        )
        if sql_match:
            return sql_match.group(0).strip()

        return text.strip()

    def text2sql(self, question: str) -> str:
        """将自然语言转换为SQL"""
        # 获取问题嵌入
        q_emb = self.local_embedding_function([question])[0]

        # RAG检索DDL
        ddl_hits = self.retrieve("ddl_knowledge", q_emb, top_k=3, fields=["ddl_text"])
        ddl_context = "\n".join(hit.get("ddl_text", "") for hit in ddl_hits)

        # RAG检索示例对
        q2sql_hits = self.retrieve(
            "q2sql_knowledge", q_emb, top_k=3, fields=["question", "sql_text"]
        )
        example_context = "\n".join(
            f"NL: \"{hit.get('question', '')}\"\nSQL: \"{hit.get('sql_text', '')}\""
            for hit in q2sql_hits
        )

        # RAG检索字段描述
        desc_hits = self.retrieve(
            "dbdesc_knowledge",
            q_emb,
            top_k=5,
            fields=["table_name", "column_name", "description"],
        )
        desc_context = "\n".join(
            f"{hit.get('table_name', '')}.{hit.get('column_name', '')}: {hit.get('description', '')}"
            for hit in desc_hits
        )

        # 如果检索失败，使用本地数据作为示例
        if not example_context and self.q2sql_pairs:
            # 使用前3个示例作为fallback
            example_context = "\n".join(
                f"NL: \"{pair['question']}\"\nSQL: \"{pair['sql']}\""
                for pair in self.q2sql_pairs[:3]
            )

        # 组装prompt
        prompt = (
            f"### Schema Definitions:\n{ddl_context}\n"
            f"### Field Descriptions:\n{desc_context}\n"
            f"### Examples:\n{example_context}\n"
            f'### Query:\n"{question}"\n'
            "请只返回SQL语句，不要包含任何解释或说明。"
        )

        try:
            response = self.openai_client.chat.completions.create(
                model="qwen-plus",
                messages=[{"role": "user", "content": prompt}],
                stream=False,
            )
            raw_sql = response.choices[0].message.content.strip()
            sql = self.extract_sql(raw_sql)
            return sql
        except Exception as e:
            logger.error(f"SQL生成失败: {e}")
            return "SELECT 1; -- Error generating SQL"

    class LocalEmbeddings:
        """本地Embedding包装器，适配RAGAS"""

        def __init__(self, embedding_function):
            self.embedding_function = embedding_function

        def embed_documents(self, texts):
            return self.embedding_function(texts)

        def embed_query(self, text):
            return self.embedding_function([text])[0]

        async def aembed_documents(self, texts):
            return self.embed_documents(texts)

        async def aembed_query(self, text):
            return self.embed_query(text)

    def evaluate_text2sql_relevancy(self, sample_size: int = 10):
        """评估Text2SQL的答案相关性"""

        # 从Q2SQL数据集中选择样本进行评估
        sample_pairs = self.q2sql_pairs[:sample_size]

        questions = []
        generated_answers = []
        reference_sqls = []
        contexts = []

        print(f"\n=== 开始Text2SQL答案相关性评估 (样本数: {len(sample_pairs)}) ===")

        results = []
        for i, pair in enumerate(sample_pairs):
            question = pair["question"]
            reference_sql = pair["sql"]

            print(f"\n处理问题 {i+1}/{len(sample_pairs)}: {question}")

            # 生成SQL
            generated_sql = self.text2sql(question)

            questions.append(question)
            generated_answers.append(generated_sql)
            reference_sqls.append(reference_sql)

            # 创建上下文（包含参考SQL作为理想答案的上下文）
            context = [
                f"The question asks for: {question}",
                f"Expected SQL format: {reference_sql}",
                "This is a database query that should return relevant results.",
            ]
            contexts.append(context)

            print(f"  生成的SQL: {generated_sql}")
            print(f"  参考SQL:   {reference_sql}")

            # 存储结果用于表格显示
            results.append(
                {
                    "问题编号": i + 1,
                    "问题": question,
                    "生成SQL": generated_sql,
                    "参考SQL": reference_sql,
                    "SQL匹配度": self._calculate_sql_match_status(
                        generated_sql, reference_sql
                    ),
                }
            )

        # 准备RAGAS评估数据
        data = {
            "question": questions,
            "answer": generated_answers,
            "contexts": contexts,
        }

        dataset = Dataset.from_dict(data)

        # 创建embedding包装器
        embedding_wrapper = self.LocalEmbeddings(self.local_embedding_function)

        # 创建答案相关性评估指标
        relevancy_metric = [AnswerRelevancy(llm=self.llm, embeddings=embedding_wrapper)]

        print(f"\n=== 开始RAGAS答案相关性评估 ===")

        try:
            result = evaluate(dataset, relevancy_metric)
            scores = result["answer_relevancy"]
            mean_score = (
                np.mean(scores) if isinstance(scores, (list, np.ndarray)) else scores
            )

            print(f"\n=== Text2SQL答案相关性评估结果 ===")
            print(f"平均相关性评分: {mean_score:.4f}")
            print(f"各问题评分: {scores}")

            # 添加评分到结果列表
            for i, score in enumerate(scores):
                results[i]["相关性评分"] = f"{score:.4f}"

            # 输出Markdown表格
            self._print_markdown_table(results, "Text2SQL答案相关性评估")

            # 详细结果分析
            print(f"\n=== 详细结果分析 ===")
            for i, (question, generated, reference, score) in enumerate(
                zip(questions, generated_answers, reference_sqls, scores)
            ):
                print(f"\n问题 {i+1}: {question}")
                print(f"  生成SQL: {generated}")
                print(f"  参考SQL: {reference}")
                print(f"  相关性评分: {score:.4f}")

                # 简单的SQL匹配度分析
                generated_clean = generated.lower().strip()
                reference_clean = reference.lower().strip()
                if generated_clean == reference_clean:
                    match_status = "完全匹配"
                elif (
                    generated_clean in reference_clean
                    or reference_clean in generated_clean
                ):
                    match_status = "部分匹配"
                else:
                    match_status = "不匹配"
                print(f"  SQL匹配度: {match_status}")

            return {
                "mean_score": mean_score,
                "individual_scores": scores,
                "questions": questions,
                "generated_answers": generated_answers,
                "reference_answers": reference_sqls,
            }

        except Exception as e:
            logger.error(f"RAGAS评估失败: {e}")
            return None

    def _calculate_sql_match_status(self, generated_sql, reference_sql):
        """计算SQL匹配状态"""
        generated_clean = generated_sql.lower().strip()
        reference_clean = reference_sql.lower().strip()
        if generated_clean == reference_clean:
            return "完全匹配"
        elif generated_clean in reference_clean or reference_clean in generated_clean:
            return "部分匹配"
        else:
            return "不匹配"

    def _print_markdown_table(self, results, title):
        """打印Markdown表格"""
        print(f"\n### {title}")
        print("| 问题编号 | 问题 | 生成SQL | 参考SQL | SQL匹配度 | 相关性评分 |")
        print("|---------|------|---------|---------|----------|-----------|")

        for result in results:
            # 处理可能包含的换行符和特殊字符
            question = result["问题"].replace("|", "\\|").replace("\n", " ")
            generated_sql = result["生成SQL"].replace("|", "\\|").replace("\n", " ")
            reference_sql = result["参考SQL"].replace("|", "\\|").replace("\n", " ")

            # 为更好的对比显示，将SQL分别放在两行
            print(
                f"| {result['问题编号']} | {question} | {generated_sql} | {reference_sql} | {result['SQL匹配度']} | {result.get('相关性评分', 'N/A')} |"
            )

    def evaluate_text2sql_faithfulness_and_relevancy(self, sample_size: int = 10):
        """评估Text2SQL的忠实度和答案相关性"""

        # 从Q2SQL数据集中选择样本进行评估
        sample_pairs = self.q2sql_pairs[:sample_size]

        questions = []
        generated_answers = []
        reference_sqls = []
        contexts = []

        print(
            f"\n=== 开始Text2SQL忠实度和答案相关性评估 (样本数: {len(sample_pairs)}) ==="
        )

        results = []
        for i, pair in enumerate(sample_pairs):
            question = pair["question"]
            reference_sql = pair["sql"]

            print(f"\n处理问题 {i+1}/{len(sample_pairs)}: {question}")

            # 生成SQL，并获取检索的上下文信息
            q_emb = self.local_embedding_function([question])[0]

            # RAG检索DDL
            ddl_hits = self.retrieve(
                "ddl_knowledge", q_emb, top_k=3, fields=["ddl_text"]
            )
            ddl_context = "\n".join(hit.get("ddl_text", "") for hit in ddl_hits)

            # RAG检索示例对
            q2sql_hits = self.retrieve(
                "q2sql_knowledge", q_emb, top_k=3, fields=["question", "sql_text"]
            )
            example_context = "\n".join(
                f"NL: \"{hit.get('question', '')}\"\nSQL: \"{hit.get('sql_text', '')}\""
                for hit in q2sql_hits
            )

            # RAG检索字段描述
            desc_hits = self.retrieve(
                "dbdesc_knowledge",
                q_emb,
                top_k=5,
                fields=["table_name", "column_name", "description"],
            )
            desc_context = "\n".join(
                f"{hit.get('table_name', '')}.{hit.get('column_name', '')}: {hit.get('description', '')}"
                for hit in desc_hits
            )

            # 生成SQL
            generated_sql = self.text2sql(question)

            questions.append(question)
            generated_answers.append(generated_sql)
            reference_sqls.append(reference_sql)

            # 创建丰富的上下文信息（用于忠实度评估）
            context = []
            if ddl_context:
                context.append(f"Database Schema: {ddl_context}")
            if desc_context:
                context.append(f"Field Descriptions: {desc_context}")
            if example_context:
                context.append(f"Similar Examples: {example_context}")
            else:
                # 如果检索失败，使用本地数据作为fallback
                fallback_examples = "\n".join(
                    f"NL: \"{pair['question']}\"\nSQL: \"{pair['sql']}\""
                    for pair in self.q2sql_pairs[:3]
                )
                context.append(f"Similar Examples: {fallback_examples}")

            context.append(f"The question asks for: {question}")
            context.append(
                f"Expected SQL format should follow database schema and return relevant results."
            )

            contexts.append(context)

            print(f"  生成的SQL: {generated_sql}")
            print(f"  参考SQL:   {reference_sql}")

            # 存储结果用于表格显示
            results.append(
                {
                    "问题编号": i + 1,
                    "问题": question,
                    "生成SQL": generated_sql,
                    "参考SQL": reference_sql,
                    "SQL匹配度": self._calculate_sql_match_status(
                        generated_sql, reference_sql
                    ),
                }
            )

        # 准备RAGAS评估数据
        data = {
            "question": questions,
            "answer": generated_answers,
            "contexts": contexts,
        }

        dataset = Dataset.from_dict(data)

        # 创建embedding包装器
        embedding_wrapper = self.LocalEmbeddings(self.local_embedding_function)

        # 创建评估指标（忠实度和答案相关性）
        metrics = [
            Faithfulness(llm=self.llm),  # 忠实度评估只需要LLM
            AnswerRelevancy(
                llm=self.llm, embeddings=embedding_wrapper
            ),  # 答案相关性需要LLM和embedding
        ]

        print(f"\n=== 开始RAGAS忠实度和答案相关性评估 ===")

        try:
            result = evaluate(dataset, metrics)

            faithfulness_scores = result["faithfulness"]
            relevancy_scores = result["answer_relevancy"]

            faithfulness_mean = (
                np.mean(faithfulness_scores)
                if isinstance(faithfulness_scores, (list, np.ndarray))
                else faithfulness_scores
            )
            relevancy_mean = (
                np.mean(relevancy_scores)
                if isinstance(relevancy_scores, (list, np.ndarray))
                else relevancy_scores
            )

            print(f"\n=== Text2SQL评估结果 ===")
            print(f"平均忠实度评分: {faithfulness_mean:.4f}")
            print(f"平均相关性评分: {relevancy_mean:.4f}")
            print(f"忠实度评分详情: {faithfulness_scores}")
            print(f"相关性评分详情: {relevancy_scores}")

            # 添加评分到结果列表
            for i, (f_score, r_score) in enumerate(
                zip(faithfulness_scores, relevancy_scores)
            ):
                results[i]["忠实度评分"] = f"{f_score:.4f}"
                results[i]["相关性评分"] = f"{r_score:.4f}"

            # 输出Markdown表格
            self._print_markdown_table(results, "Text2SQL忠实度和答案相关性评估")

            # 详细结果分析
            print(f"\n=== 详细结果分析 ===")
            for i, (question, generated, reference, f_score, r_score) in enumerate(
                zip(
                    questions,
                    generated_answers,
                    reference_sqls,
                    faithfulness_scores,
                    relevancy_scores,
                )
            ):
                print(f"\n问题 {i+1}: {question}")
                print(f"  生成SQL: {generated}")
                print(f"  参考SQL: {reference}")
                print(f"  忠实度评分: {f_score:.4f}")
                print(f"  相关性评分: {r_score:.4f}")

                # 简单的SQL匹配度分析
                generated_clean = generated.lower().strip()
                reference_clean = reference.lower().strip()
                if generated_clean == reference_clean:
                    match_status = "完全匹配"
                elif (
                    generated_clean in reference_clean
                    or reference_clean in generated_clean
                ):
                    match_status = "部分匹配"
                else:
                    match_status = "不匹配"
                print(f"  SQL匹配度: {match_status}")

            return {
                "faithfulness_mean": faithfulness_mean,
                "relevancy_mean": relevancy_mean,
                "faithfulness_scores": faithfulness_scores,
                "relevancy_scores": relevancy_scores,
                "questions": questions,
                "generated_answers": generated_answers,
                "reference_answers": reference_sqls,
            }

        except Exception as e:
            logger.error(f"RAGAS评估失败: {e}")
            return None

    def run_interactive_demo(self):
        """运行交互式演示"""
        print("\n=== Text2SQL交互式演示 ===")
        print("输入自然语言问题，我将生成SQL并评估相关性")
        print("输入 'quit' 或 'exit' 退出")

        while True:
            question = input("\n请输入问题: ").strip()

            if question.lower() in ["quit", "exit", "退出"]:
                break

            if not question:
                continue

            print(f"\n正在处理: {question}")

            # 生成SQL
            generated_sql = self.text2sql(question)
            print(f"生成的SQL: {generated_sql}")

            # 单独评估这个问题的相关性
            data = {
                "question": [question],
                "answer": [generated_sql],
                "contexts": [
                    [
                        f"The question asks for: {question}",
                        "This is a database query that should return relevant results.",
                        "The SQL should be syntactically correct and semantically meaningful.",
                    ]
                ],
            }

            dataset = Dataset.from_dict(data)
            embedding_wrapper = self.LocalEmbeddings(self.local_embedding_function)
            relevancy_metric = [
                AnswerRelevancy(llm=self.llm, embeddings=embedding_wrapper)
            ]

            try:
                result = evaluate(dataset, relevancy_metric)
                score = result["answer_relevancy"][0]
                print(f"答案相关性评分: {score:.4f}")

                faithfulness_metric = [Faithfulness(llm=self.llm)]
                faithfulness_result = evaluate(dataset, faithfulness_metric)
                f_score = faithfulness_result["faithfulness"][0]
                print(f"忠实度评分: {f_score:.4f}")

            except Exception as e:
                print(f"评估失败: {e}")


def main():
    """主函数"""
    evaluator = Text2SQLRAGASEvaluator()

    print("选择运行模式:")
    print("1. 批量评估 - 仅答案相关性 (使用Q2SQL数据集)")
    print("2. 批量评估 - 忠实度和答案相关性 (使用Q2SQL数据集)")
    print("3. 交互式演示")

    choice = input("请选择 (1, 2或 3): ").strip()

    if choice == "1":
        sample_size = input("请输入评估样本数量 (默认10): ").strip()
        sample_size = int(sample_size) if sample_size.isdigit() else 10
        evaluator.evaluate_text2sql_relevancy(sample_size)
    elif choice == "2":
        sample_size = input("请输入评估样本数量 (默认10): ").strip()
        sample_size = int(sample_size) if sample_size.isdigit() else 10
        evaluator.evaluate_text2sql_faithfulness_and_relevancy(sample_size)
    elif choice == "3":
        evaluator.run_interactive_demo()
    else:
        print("无效选择")


if __name__ == "__main__":
    main()
